package com.bibimaya.auto.mapper.dao.scan;

import java.io.File;
import java.io.FileFilter;
import java.io.IOException;
import java.net.JarURLConnection;
import java.net.URL;
import java.net.URLDecoder;
import java.util.Enumeration;
import java.util.LinkedHashSet;
import java.util.Set;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

/**
 * Class 扫描类
 */
public class ClassScaner {


    public static Set<Class<?>> getClass(String basePackage, boolean recursive) {
        Set<Class<?>> classes = new LinkedHashSet<Class<?>>();
        String packageName = basePackage;
        if (packageName.endsWith(".")) {
            packageName = packageName.substring(0, packageName.lastIndexOf('.'));
        }
        String package2Path = packageName.replace('.', '/');

        Enumeration<URL> dirs;
        try {
            dirs = Thread.currentThread().getContextClassLoader().getResources(package2Path);
            while (dirs.hasMoreElements()) {
                URL url = dirs.nextElement();
                String protocol = url.getProtocol();
                if ("file".equals(protocol)) {
                    String filePath = URLDecoder.decode(url.getFile(), "UTF-8");
                    scanClassesByFile(classes, packageName, filePath, recursive);
                } else if ("jar".equals(protocol)) {
                    scanClassesByJar(packageName, url, recursive, classes);
                }
            }
        } catch (IOException e) {
            e.printStackTrace();
        }

        return classes;
    }

    /**
     * 以文件的方式扫描包下的所有Class文件
     * @param packageName
     * @param packagePath
     * @param recursive
     * @param classes
     */
    private static void scanClassesByFile(Set<Class<?>> classes, String packageName, String packagePath,
                                            boolean recursive) {
        File dir = new File(packagePath);
        if (!dir.exists() || !dir.isDirectory()) {
            return;
        }
        final boolean fileRecursive = recursive;
        File[] dirfiles = dir.listFiles(new FileFilter() {
            @Override
            public boolean accept(File file) {
                if (file.isDirectory()) {
                    return fileRecursive;
                }
                String filename = file.getName();
                if (filename.indexOf('$') != -1) {
                    return false;
                }
                return true;
            }
        });
        for (File file : dirfiles) {
            if (file.isDirectory()) {
                scanClassesByFile(classes, packageName + "." + file.getName(), file.getAbsolutePath(), recursive);
            } else {
                String className = file.getName().substring(0, file.getName().length() - 6);
                try {
                    classes.add(Thread.currentThread().getContextClassLoader().loadClass(packageName + '.' + className));

                } catch (ClassNotFoundException e) {
                    e.printStackTrace();
                }
            }
        }
    }


    private static void scanClassesByJar(String basePackage, URL url, final boolean recursive, Set<Class<?>> classes) {
        String packageName = basePackage;
        String package2Path = packageName.replace('.', '/');
        JarFile jar;
        try {
            jar = ((JarURLConnection) url.openConnection()).getJarFile();
            Enumeration<JarEntry> entries = jar.entries();
            while (entries.hasMoreElements()) {
                JarEntry entry = entries.nextElement();
                String name = entry.getName();
                if (!name.startsWith(package2Path) || entry.isDirectory()) {
                    continue;
                }

                // 判断是否递归搜索子包
                if (!recursive && name.lastIndexOf('/') != package2Path.length()) {
                    continue;
                }
                // 判断是否过滤 inner class
                if (name.indexOf('$') != -1) {
                    continue;
                }
                String className = name.replace('/', '.');
                className = className.substring(0, className.length() - 6);
                classes.add(Thread.currentThread().getContextClassLoader().loadClass(className));
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static void main(String[] args){

        Set<Class<?>> allClass = ClassScaner.getClass("com.bibimaya",true);
        for(Class c : allClass){
            System.out.println(c.getCanonicalName());
        }
    }

}
